package com.example.demo.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.example.demo.Mapper.DocumentMapper;
import com.example.demo.Mapper.MessageMapper;
import com.example.demo.Mapper.SessionMapper;
import com.example.demo.context.BaseContext;
import com.example.demo.entity.Document;
import com.example.demo.entity.Message;
import com.example.demo.entity.Session;
import com.example.demo.entity.StreamEventType;
import com.example.demo.service.RagAssistant;
import com.example.demo.service.RagSerivce;
import com.example.demo.service.SummaryAssistant;
import com.example.demo.vo.documentDetailVO;
import com.example.demo.vo.messageDetailVO;
import dev.langchain4j.data.message.ChatMessageType;
import lombok.RequiredArgsConstructor;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.time.LocalDateTime;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;

@Service
@RequiredArgsConstructor
public class RagServiceImpl extends ServiceImpl<DocumentMapper, Document> implements RagSerivce {
    final DocumentMapper documentMapper;

    final RagAssistant assistant;

    final SummaryAssistant summaryAssistant;

    final SessionMapper sessionMapper;

    final MessageMapper messageMapper;

    @Override
    public List<documentDetailVO> urlList(Long userId) {
        return documentMapper.getDocumentDetail(userId);
    }


    @Override
    public List<messageDetailVO> getMessageDetail(Long userId) {
        return messageMapper.getRagMessage(userId,1);
    }

    @Override
    public Flux<ServerSentEvent<Object>> streamResponse(String message, Long sessionId) {
        Long userId= BaseContext.getCurrentId();
        // 1.将用户发送的信息存入数据库
        Message userMessage = new Message();
        userMessage.setUserId(userId);
        userMessage.setSessionId(sessionId);
        userMessage.setContent(message);
        userMessage.setTag(1);
        userMessage.setType(ChatMessageType.USER);
        userMessage.setCreatedAt(LocalDateTime.now());
        messageMapper.insert(userMessage);

        // 2. 检查是否需要发送标题事件
        Mono<ServerSentEvent<Object>> titleEvent = Mono.fromCallable(() -> {
                    // 查询会话是否有标题
                    String currentTitle = sessionMapper.getSessionTitle(sessionId);
                    if (currentTitle == null || currentTitle.isEmpty()) {
                        // 从用户消息生成标题（前20个字符或完整消息）
                        //让AI总结这个会话的标题
                        String generatedTitle = summaryAssistant.chat(message);

                        // 更新会话标题
                        sessionMapper.updateSessionTitle(sessionId, generatedTitle);

                        return ServerSentEvent.builder()
                                .event(StreamEventType.TITLE.getValue())
                                .data(Map.of("title", generatedTitle))
                                .build();
                    }
                    return null; // 不需要发送标题事件
                })
                .filter(Objects::nonNull) // 过滤掉null事件
                .subscribeOn(Schedulers.boundedElastic()); // 异步执行

        // 2. 初始化事件
        Mono<ServerSentEvent<Object>> readyEvent = Mono.just(
                ServerSentEvent.builder()
                        .event(StreamEventType.READY.getValue())
                        .data(Collections.emptyMap())
                        .build()
        );

        // 3. 初始化空消息结构
        Mono<ServerSentEvent<Object>> initEvent = Mono.just(
                ServerSentEvent.builder()
                        .data(messageDetailVO.builder()
                                .userId(BaseContext.getCurrentId())  // 实际应从安全上下文获取
                                .sessionId(sessionId)
                                .content("")
                                .type(ChatMessageType.AI)
                                .createdAt(LocalDateTime.now())
                                .build())
                        .build()
        );

        // 用于收集AI回复的所有内容
        StringBuilder aiResponseContent = new StringBuilder();


        // 4. 流式内容处理（核心修改部分）
        Flux<ServerSentEvent<Object>> contentStream = assistant.chat(sessionId.toString(),message)
                .doOnNext(aiResponseContent::append)  // 收集所有内容块
                .onErrorResume(e -> {
                    // 记录错误日志
                    log.error("Stream processing error", e);
                    // 可以选择在这里保存部分响应
                    return Mono.empty();
                })
                .map(chunk -> ServerSentEvent.builder()
                        .data(Map.of(
                                "path", "content",
                                "operation", "APPEND",
                                "value", chunk
                        ))
                        .build()
                );


        // 5. 结束事件，并在结束时保存AI消息到数据库
        Mono<ServerSentEvent<Object>> finishEvent = Mono.defer(() -> {
            // 创建并保存AI消息
            Message aiMessage = new Message();
            aiMessage.setUserId(userId);
            aiMessage.setSessionId(sessionId);
            aiMessage.setContent(aiResponseContent.toString());
            aiMessage.setTag(1);
            aiMessage.setType(ChatMessageType.AI);
            aiMessage.setCreatedAt(LocalDateTime.now());

            return Mono.fromRunnable(() -> messageMapper.insert(aiMessage))
                    .subscribeOn(Schedulers.boundedElastic())  // 在单独线程执行
                    .thenReturn(ServerSentEvent.builder().event(StreamEventType.FINISH.getValue())
                            .data(Collections.emptyMap())
                            .build());
        });

        // 组合事件流：ready -> title(如果有) -> init -> content -> finish
        return Flux.concat(
                readyEvent,
                titleEvent.defaultIfEmpty(ServerSentEvent.builder().build()), // 空事件占位
                initEvent,
                contentStream,
                finishEvent
        );
    }
}
